"""
主题: 实现消息发布/订阅模型
问题: 你有一个基于线程通信的程序，想让它们实现发布/订阅模式的消息通信。
提示 : 
"""
from contextlib import contextmanager
from collections import defaultdict

'''简单的交换机实现'''

class Exchange:
    def __init__(self):
        self._subscribers = set()

    def attach(self, task):
        self._subscribers.add(task)

    def detach(self, task):
        self._subscribers.remove(task)

    def send(self, msg):
        for subscriber in self._subscribers:
            subscriber.send(msg)

    @contextmanager
    def subscribe(self, *tasks):
        for task in tasks:
            self.attach(task)
        try:
            yield
        finally:
            for task in tasks:
                self.detach(task)

# Dictionary of all created exchanges
_exchanges = defaultdict(Exchange)

# Return the Exchange instance associated with a given name
def get_exchange(name):
    return _exchanges[name]

class DisplayMessages:
    '''简单的诊断类'''
    def __init__(self):
        self.count = 0
    def send(self, msg):
        self.count += 1
        print('msg[{}]: {!r}'.format(self.count, msg))

    # Example task (just for testing)
class Task:
    def __init__(self, name):
        self.name = name
    def send(self, msg):
        print('{} got: {!r}'.format(self.name, msg))

def recipe1():
    task_a = Task('A')
    task_b = Task('B')

    exc = get_exchange('spam')
    exc.attach(task_a)
    exc.attach(task_b)
    d = DisplayMessages()
    exc.attach(d)
    exc.send('msg1')
    exc.send('msg2')

    exc.detach(task_a)
    exc.detach(task_b)
    exc.send('msg3')

def recipe2():
    task_a = Task('A')
    task_b = Task('B')

    exc = get_exchange('spam')
    with exc.subscribe(task_a, task_b):
        exc.send('msg1')
        exc.send('msg2')

    exc.send('msg3')


def main():
    print('recipe1'.center(20, '*'))
    recipe1()
    print('recipe2'.center(20, '*'))
    recipe2()


if __name__ == '__main__':
    main()   



